#include "MyMath.h"

using namespace Eigen;

int MatrixSolver::SolveMatrix(std::vector<double> &rhs, std::vector<double> &solution)
{
	int vertexNumber = rhs.size();//the two vectors have the same size
	//Map is used because internal type is required when solving the matrix
	Eigen::Map<VectorXd> bToMap(rhs.data(), vertexNumber, 1);
	Eigen::Map<VectorXd> xToMap(solution.data(), vertexNumber, 1);

	//use SparseLU to solve sparse matrix problem. SparseLU supports general square sparse systems
	Eigen::SparseLU<Eigen::SparseMatrix<double>,Eigen::COLAMDOrdering<int>> theSparseSolver;
	theSparseSolver.analyzePattern(this->sparseMatrix);
	theSparseSolver.factorize(sparseMatrix);
	xToMap = theSparseSolver.solve(bToMap);

	if (theSparseSolver.info() != Eigen::Success)
	{
		std::cerr << "Solve Ax=b failed in this right-hand side" << std::endl;
		return 1;
	}
	return 0;
}

OneDimSolver::OneDimSolver(int vertexNum)
{
	this->oneDimSize = vertexNum;
}

int OneDimSolver::SolveOneDim(std::vector<double>& rhs, std::vector<double>& x, 
						const std::vector<double> &eps, const std::vector<double> &deltaX, const OneDimBC &left, const OneDimBC &right)
{
	buildSparseMatrix(eps, deltaX, left, right);
	//use to examine the sparse matrix 
	//std::cout << sparseMatrix << std::endl;
	int ret = 0;
	ret = SolveMatrix(rhs, x);
	
	return ret;
}

int OneDimSolver::buildSparseMatrix(const std::vector<double> &eps, const std::vector<double> &deltaX, const OneDimBC &left, const OneDimBC &right)
{
	sparseMatrix.resize(oneDimSize, oneDimSize);
	sparseMatrix.reserve(Eigen::VectorXd::Constant(oneDimSize,3));//reserve room for 3 non-zeros per column
	
	for (int i=1; i<oneDimSize-1; ++i)
	{
		sparseMatrix.insert(i, i-1) = eps[i-1] / deltaX[i-1];
		sparseMatrix.insert(i, i)   = - (eps[i-1] / deltaX[i-1] + eps[i] / deltaX[i]);
		sparseMatrix.insert(i, i+1) = eps[i] / deltaX[i];
	}
	//build the left boundary condition
	if (left.Type == Utilities::First)
	{
		sparseMatrix.insert(0, 0) = 1;
	}
	else if (left.Type == Utilities::Second)
	{
		sparseMatrix.insert(0, 0) = - eps[0] / deltaX[0];
		sparseMatrix.insert(0, 1) = eps[0] / deltaX[0];
	}
	//build the right boundary condition
	int lastElementIndex = eps.size() - 1;
	if (right.Type == Utilities::First)
	{
		sparseMatrix.insert(oneDimSize-1, oneDimSize-1) = 1;
	}
	else if (right.Type == Utilities::Second)
	{
		//not that the last element of deltaX is deltaX[oneDimSize-2]
		sparseMatrix.insert(oneDimSize-1, oneDimSize-2) = eps[lastElementIndex] / deltaX[lastElementIndex];
		sparseMatrix.insert(oneDimSize-1, oneDimSize-1) = - eps[lastElementIndex] / deltaX[lastElementIndex];
	}
	sparseMatrix.makeCompressed();
	return 0;
}

TwoDimSolver::TwoDimSolver(int vertexNumX, int vertexNumY)
{
	this->twoDimSizeX = vertexNumX;
	this->twoDimSizeY = vertexNumY;
	this->twoDimSize = twoDimSizeX * twoDimSizeY;
}

int TwoDimSolver::buildSparseMatrix(const std::vector<double> &meshX, const std::vector<double> &meshY, const My2DVector<double> &eps)
{
	sparseMatrix.resize(twoDimSize, twoDimSize);
	sparseMatrix.reserve(Eigen::VectorXd::Constant(twoDimSize, 5));//reserver room for 5 non-zeros per column

	My2DVector<double> positioner = My2DVector<double>(twoDimSizeX, twoDimSizeY);
	int row = 0;
	int col = 0;
	for (int iy = 0; iy != twoDimSizeX; ++iy)
	{
		for (int ix = 0; ix != twoDimSizeY; ++ix)
		{
			row = positioner.IndexAt(ix, iy);
			
			//left
			if (ix != 0)
			{
				col = positioner.IndexAt(ix-1, iy);
				sparseMatrix.insert(row, col) = 0.5 * ( ((ix==0             || iy==twoDimSizeY-1) ? 0 : meshY[iy]   / meshX[ix-1] * eps(ix-1, iy)) //left up
													   +((ix==0             || iy==0)             ? 0 : meshY[iy-1] / meshX[ix-1] * eps(ix-1, iy-1))); //left down
			}
			//right
			if (ix != twoDimSizeX-1)
			{
				col = positioner.IndexAt(ix+1, iy);
				sparseMatrix.insert(row, col) = 0.5 * ( ((ix==twoDimSizeX-1 || iy==0)             ? 0 : meshY[iy-1] / meshX[ix]   * eps(ix, iy-1)) //right down
													   +((ix==twoDimSizeX-1 || iy==twoDimSizeY-1) ? 0 : meshY[iy]   / meshX[ix]   * eps(ix, iy))); //right up
			}
			

			//center
			//!important negative
			col = positioner.IndexAt(ix, iy);
			sparseMatrix.insert(row, col)   = - 0.5 * ( ((ix==0             || iy==twoDimSizeY-1) ? 0 : meshY[iy]   / meshX[ix-1] * eps(ix-1, iy)   +  meshX[ix-1] / meshY[iy]   * eps(ix-1, iy)) //left up
													   +((ix==0             || iy==0)             ? 0 : meshY[iy-1] / meshX[ix-1] * eps(ix-1, iy-1) +  meshX[ix-1] / meshY[iy-1] * eps(ix-1, iy-1)) //left down
													   +((ix==twoDimSizeX-1 || iy==twoDimSizeY-1) ? 0 : meshY[iy]   / meshX[ix]   * eps(ix, iy)     +  meshX[ix]   / meshY[iy]   * eps(ix, iy)) //right up
													   +((ix==twoDimSizeX-1 || iy==0			) ? 0 : meshX[ix]   / meshY[iy-1] * eps(ix, iy-1)   +  meshY[iy-1] / meshX[ix]   * eps(ix, iy-1))); //right down

			//down
			if (iy != 0)
			{
				col = positioner.IndexAt(ix, iy-1);
				sparseMatrix.insert(row, col) = 0.5 * ( ((ix==0             || iy==0)             ? 0 : meshX[ix-1] / meshY[iy-1] * eps(ix-1, iy-1)) //left down
													   +((ix==twoDimSizeX-1 || iy==0)             ? 0 : meshX[ix]   / meshY[iy-1] * eps(ix, iy-1))); //right down
			}

			
			//up
			if (iy != twoDimSizeY-1)
			{			
				col = positioner.IndexAt(ix, iy+1);
				sparseMatrix.insert(row, col) = 0.5 * ( ((ix==twoDimSizeX-1 || iy==twoDimSizeY-1) ? 0 : meshX[ix]   / meshY[iy]   * eps(ix, iy)) //right up
													   +((ix==0             || iy==twoDimSizeY-1) ? 0 : meshX[ix-1] / meshY[iy]   * eps(ix-1, iy))); //right down
			}
		}
	}
	return 0;
}

int TwoDimSolver::SolveTwoDim(std::vector<double> &rhs, std::vector<double> &x, 
							  const std::vector<double> &meshX, const std::vector<double> &meshY, const My2DVector<double> &eps, 
							  const TwoDimBC &left, const TwoDimBC &right, const TwoDimBC &up, const TwoDimBC &down)
{
	buildSparseMatrix(meshX, meshY, eps);
	refreshSparseMatrix(left, right, up, down);
	//data monitoring
	//std::cout << sparseMatrix << std::endl;
	int ret = 0;
	ret = this->SolveMatrix(rhs, x);
	return ret;
}

int TwoDimSolver::refreshSparseMatrix(const TwoDimBC &left, const TwoDimBC &right, const TwoDimBC &up, const TwoDimBC &down)
{
	My2DVector<double> positioner = My2DVector<double>(twoDimSizeX, twoDimSizeY);
	int ix = 0;
	int iy = 0;
	int index = 0;

	//---------------------------------------------------------------------------
	//the follow method is used to iterate the non-zero coefficient of the matrix
	//for (int k=0; k<sparseMatrix.outerSize(); ++k)
	//	for (SparseMatrix<double>::InnerIterator it(sparseMatrix,k); it; ++it)
	//	{
	//		it.valueRef() = 9; // for get the reference of the coefficient
	//		it.row(); // row index
	//		it.col(); // col index (here it is equal to k)
	//		it.index(); // inner index, here it is equal to it.row()
	//	}
	//----------------------------------------------------------------------------

	//if the corner point is included to Dirichlet boundary conditions of one boundary, it becomes a point with known potential
	//i.e. Dirichlet boundary condition is prior to Neumann boundary condition
	//only refresh the matrix when it is Dirichlet boundary conditions
	switch (left.Type)//left boundary
	{
	case First:
		for (int k=0; k<sparseMatrix.outerSize(); ++k)
			for (SparseMatrix<double>::InnerIterator it(sparseMatrix,k); it; ++it)
			{
				index = it.row();
				ix = positioner.GetX(index);

				if (ix == 0)//find the point on the left boundary
				{
					if (it.row() == it.col())
						it.valueRef() = 1;
					else
						it.valueRef() = 0;
				}
			}
		break;
	}

	switch (right.Type)//right boundary
	{
	case First:
		for (int k=0; k<sparseMatrix.outerSize(); ++k)
			for (SparseMatrix<double>::InnerIterator it(sparseMatrix,k); it; ++it)
			{
				index = it.row();
				ix = positioner.GetX(index);

				if (ix == twoDimSizeX-1)//find the point on the right boundary
				{
					if (it.row() == it.col())
						it.valueRef() = 1;
					else
						it.valueRef() = 0;
				}
			}
		break;
	}

	switch (up.Type)
	{
	case First:
		for (int k=0; k<sparseMatrix.outerSize(); ++k)
			for (SparseMatrix<double>::InnerIterator it(sparseMatrix,k); it; ++it)
			{
				index = it.row();
				iy = positioner.GetY(index);

				if (iy == twoDimSizeY-1)//find the point on the up boundary
				{
					if (it.row() == it.col())
						it.valueRef() = 1;
					else
						it.valueRef() = 0;
				}
			}
		break;
	}

	switch (down.Type)
	{
	case First:
		for (int k=0; k<sparseMatrix.outerSize(); ++k)
			for (SparseMatrix<double>::InnerIterator it(sparseMatrix,k); it; ++it)
			{
				index = it.row();
				iy = positioner.GetY(index);

				if (iy == 0)//find the point on the down boundary
				{
					if (it.row() == it.col())
						it.valueRef() = 1;
					else
						it.valueRef() = 0;
				}
			}
		break;
	}
	sparseMatrix.makeCompressed();
	//for data monitoring
	//std::cout << sparseMatrix << std::endl;
	return 0;
}
